package cn.kaiux.bigdata.homework.week05_06

import org.apache.spark.TaskContext
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.SparkSession

/**
 * @author kaiux@foxmail.com
 * @date 2022/6/9 15:50
 */
object DistCP {

  def main(args: Array[String]): Unit = {

    val options = DistCPOptionsParsing.parse(args)

    val sparkSession = SparkSession
      .builder()
      .appName("Spark DistCP")
      .getOrCreate()

    val sourcePaths = options.getPaths._1
    val targetPath = options.getPaths._2

    val qualifiedSourcePaths = sourcePaths.map(FileUtils.pathToQualifiedPath(sparkSession.sparkContext.hadoopConfiguration, _))
    val qualifiedTargetPath = FileUtils.pathToQualifiedPath(sparkSession.sparkContext.hadoopConfiguration, targetPath)

    val sourceRDD = FileUtils.getSourceFiles(sparkSession.sparkContext, qualifiedSourcePaths.map(_.toUri), qualifiedTargetPath.toUri, options.maxConcurrence)

    val targetRDD = FileUtils.getDestinationFiles(sparkSession.sparkContext, qualifiedTargetPath, options.maxConcurrence)

    val joined = sourceRDD.fullOuterJoin(targetRDD)

    val toCopy = joined.collect { case (_, (Some(s), _)) => s }

    doCopy(toCopy, options)
      .map(_.getMessage)
      .collect()
      .foreach(println)

    sparkSession.stop()

  }

  def doCopy(sourceRDD: RDD[CopyDefinitionWithDependencies], option: Config): RDD[DistCPResult] = {
    val hadoopConfiguration = new WritableConfig(sourceRDD.sparkContext.hadoopConfiguration)
    sourceRDD
      .mapPartitions { iterator =>
        val attemptID = TaskContext.get().taskAttemptId()
        val fsCache = new FileSystemObjectCacher(hadoopConfiguration.get())
        iterator
          .flatMap(_.getAllCopyDefinitions)
          .collectMapWithEmptyCollection(
            (d, z) => z.contains(d),
            d => {
              val r = FileUtils.doCopy(fsCache.getOrCreate(d.source.uri), fsCache.getOrCreate(d.destination), d, attemptID, option.ignoreFailures)
              r
            }
          )
      }
  }

  implicit class DistCPIteratorImplicit[B](iterator: Iterator[B]) {
    def collectMapWithEmptyCollection(skip: (B, Set[B]) => Boolean, action: B => DistCPResult): Iterator[DistCPResult] = {
      iterator
        .scanLeft((Set.empty[B], None: Option[DistCPResult])) {
          case ((z, _), d) if skip(d, z) => (z, None)
          case ((z, _), d) =>
            (z + d, Some(action(d)))
        }
        .collect { case (_, Some(r)) => r }
    }
  }

}